import random
import numpy as np

# 生成一个简化版的MNIST数据集样本
def load_mnist_sample(num_samples=100):
    np.random.seed(42)
    data = []
    labels = []
    for _ in range(num_samples):
        digit = random.randint(0, 9)
        image = np.random.rand(28, 28) + digit * 0.1  # 添加一些随机噪声
        data.append(image.flatten())
        labels.append(digit)
    return np.array(data), np.array(labels)

# 计算两个向量之间的欧几里得距离
def euclidean_distance(x1, x2):
    return np.sqrt(np.sum((x1 - x2) ** 2))

# 初始化聚类中心
def initialize_centroids(data, k):
    indices = random.sample(range(len(data)), k)
    return data[indices]

# 将数据点分配到最近的聚类中心
def assign_clusters(data, centroids):
    clusters = [[] for _ in range(len(centroids))]
    for point in data:
        distances = [euclidean_distance(point, centroid) for centroid in centroids]
        cluster_index = np.argmin(distances)
        clusters[cluster_index].append(point)
    return clusters

# 更新聚类中心
def update_centroids(clusters):
    return [np.mean(cluster, axis=0) for cluster in clusters]

# k-means聚类算法
def kmeans(data, k, max_iterations=100):
    centroids = initialize_centroids(data, k)
    for _ in range(max_iterations):
        clusters = assign_clusters(data, centroids)
        new_centroids = update_centroids(clusters)
        # 修复比较逻辑
        if np.all([np.allclose(centroid, new_centroid) for centroid, new_centroid in zip(centroids, new_centroids)]):
            break
        centroids = new_centroids
    return centroids, clusters

# 主函数
if __name__ == "__main__":
    # 加载简化版的MNIST数据集
    data, labels = load_mnist_sample(num_samples=100)
    
    # 设置聚类的类别数
    k = 10
    
    # 运行k-means聚类算法
    centroids, clusters = kmeans(data, k)
    
    # 打印聚类结果
    for i, cluster in enumerate(clusters):
        print(f"Cluster {i}: {len(cluster)} points")
